import onnx
from pytorch2keras import pytorch_to_keras
import tensorflow as tf
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torchvision
import torchvision.transforms as transforms
import shutil
import numpy as np


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)  # (32, 26, 26)
        self.conv2 = nn.Conv2d(32, 64, 3)  # (64, 24, 24)
        self.pool = nn.MaxPool2d(2, 2)  # (64, 12, 12)
        self.fc1 = nn.Linear(12 * 12 * 64, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu6(self.conv1(x))
        x = self.pool(F.relu6(self.conv2(x)))
        x = x.view(-1, 12 * 12 * 64)
        x = F.relu6(self.fc1(x))
        x = self.fc2(x)
        return x


net = Net()
net.load_state_dict(torch.load('data\\mnist_net.pth'))
net.eval()
input_np = np.random.uniform(0, 1, (1, 1, 28, 28))
input_var = Variable(torch.FloatTensor(input_np))
k_model = pytorch_to_keras(net, input_var, [(1, 28, 28)], change_ordering=True)

shutil.rmtree('saved_model', ignore_errors=True)
tf.saved_model.save(k_model, 'saved_model')

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize(0.5, 0.5)])
testset = torchvision.datasets.MNIST(
    root='./data',
    train=False,
    download=True,
    transform=transform)
testloader = torch.utils.data.DataLoader(
    testset,
    batch_size=1,
    shuffle=False,
    num_workers=0)

print('TEST')
for image, label in testloader:
    image = np.ascontiguousarray(np.transpose(image.numpy(), (0, 2, 3, 1)))
    y = k_model.predict(image)[0]
    print('label={}, y={}'.format(label, y.argmax()))
